package com.weizhu.swufer.core.loadbalance;

import com.weizhu.swufer.core.entity.RpcBaseInfo;

import java.io.UnsupportedEncodingException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.TreeSet;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * 一致性哈希算法
 */
public class ConsistentHashLoadBalancer extends AbstractLoadBalancer{

    private static final int VIRTUAL_NODE_NUM = 160;

    private ConcurrentMap<String,ConsistentHashSelector> selectors = new ConcurrentHashMap<>();

    @Override
    public String doSelect(String serviceKey, List<RpcBaseInfo> rpcList, long listHashCode) {

        ConsistentHashSelector selector = selectors.get(serviceKey);
        if (selector == null || selector.listHashCode != listHashCode){
            selector = new ConsistentHashSelector(rpcList, listHashCode);
            selectors.put(serviceKey,selector);
        }
        return selector.select(serviceKey);
    }

    private static final class ConsistentHashSelector{
        long listHashCode;

        TreeMap<Long,RpcBaseInfo> treeMap = new TreeMap<>();

        ConsistentHashSelector(List<RpcBaseInfo> rpcList,long listHashCode){
            this.listHashCode = listHashCode;
            for (RpcBaseInfo rpcBaseInfo : rpcList) {
                String addressPort = rpcBaseInfo.getAddress() + rpcBaseInfo.getPort();
                for (int i = 0;i < VIRTUAL_NODE_NUM / 4;i++){
                    for (int h = 0;h < 4;h++){
                        long hashValue = hash(addressPort + i, h);
                        treeMap.put(hashValue,rpcBaseInfo);
                    }

                }
            }
        }

        private String select(String serviceKey){
            int rand = (int) (Math.random() * 4);
            long hash = hash(serviceKey + rand, rand);
            SortedMap<Long, RpcBaseInfo> last = treeMap.tailMap(hash);
            if (!last.isEmpty()) {
                RpcBaseInfo rpcBaseInfo = last.get(last.firstKey());
                return rpcBaseInfo.getAddress();
            }
            return treeMap.firstEntry().getValue().getAddress();
        }

        /**
         * get hash code on 2^32 ring (md5散列的方式计算hash值)
         * @param key
         * @return
         */
        private long hash(String key,int number) {

            // md5 byte
            MessageDigest md5;
            try {
                md5 = MessageDigest.getInstance("MD5");
            } catch (NoSuchAlgorithmException e) {
                throw new RuntimeException("MD5 not supported", e);
            }
            md5.reset();
            byte[] keyBytes = null;
            try {
                keyBytes = key.getBytes("UTF-8");
            } catch (UnsupportedEncodingException e) {
                throw new RuntimeException("Unknown string :" + key, e);
            }

            md5.update(keyBytes);
            byte[] digest = md5.digest();

            return (((long) (digest[3 + number * 4] & 0xFF) << 24)
                    | ((long) (digest[2 + number * 4] & 0xFF) << 16)
                    | ((long) (digest[1 + number * 4] & 0xFF) << 8)
                    | (digest[number * 4] & 0xFF))
                    & 0xFFFFFFFFL;

        }


    }
}
